package zzz.crypto;

/*
 * Copyright 1997-2001 by Sun Microsystems, Inc.,
 * 901 San Antonio Road, Palo Alto, California, 94303, U.S.A.
 * All rights reserved.
 *
 * This software is the confidential and proprietary information
 * of Sun Microsystems, Inc. ("Confidential Information").  You
 * shall not disclose such Confidential Information and shall use
 * it only in accordance with the terms of the license agreement
 * you entered into with Sun.
 */

import java.math.BigInteger;
import java.security.AlgorithmParameterGenerator;
import java.security.AlgorithmParameters;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PublicKey;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.SecretKey;
import javax.crypto.ShortBufferException;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;

/**
 * This program executes the Diffie-Hellman key agreement protocol between 2
 * parties: Alice and Bob.
 * 
 * By default, preconfigured parameters (1024-bit prime modulus and base
 * generator used by SKIP) are used. If this program is called with the "-gen"
 * option, a new set of parameters is created.
 */

public class DHKeyAgreement2 {

	private DHKeyAgreement2() {
	}

	public static void main(String argv[]) {
		try {
			String mode = "USE_SKIP_DH_PARAMS";

			DHKeyAgreement2 keyAgree = new DHKeyAgreement2();

			if (argv.length > 1) {
				keyAgree.usage();
				throw new Exception("Wrong number of command options");
			} else if (argv.length == 1) {
				if (!(argv[0].equals("-gen"))) {
					keyAgree.usage();
					throw new Exception("Unrecognized flag: " + argv[0]);
				}
				mode = "GENERATE_DH_PARAMS";
			}

			keyAgree.run(mode);
		} catch (Exception e) {
			System.err.println("Error: " + e);
			System.exit(1);
		}
	}

	private void run(String mode) throws Exception {

		DHParameterSpec dhSkipParamSpec;

		if (mode.equals("GENERATE_DH_PARAMS")) {
			// Some central authority creates new DH parameters
			System.out
					.println("Creating Diffie-Hellman parameters (takes VERY long) ...");
			AlgorithmParameterGenerator paramGen = AlgorithmParameterGenerator
					.getInstance("DH");
			paramGen.init(512);
			AlgorithmParameters params = paramGen.generateParameters();
			dhSkipParamSpec = (DHParameterSpec) params
					.getParameterSpec(DHParameterSpec.class);
		} else {
			// use some pre-generated, default DH parameters
			System.out.println("Using SKIP Diffie-Hellman parameters");
			dhSkipParamSpec = new DHParameterSpec(skip1024Modulus, skip1024Base);
		}

		/*
		 * Alice creates her own DH key pair, using the DH parameters from above
		 */
		System.out.println("ALICE: Generate DH keypair ...");
		KeyPairGenerator aliceKpairGen = KeyPairGenerator.getInstance("DH");
		aliceKpairGen.initialize(dhSkipParamSpec);
		KeyPair aliceKpair = aliceKpairGen.generateKeyPair();

		// Alice creates and initializes her DH KeyAgreement object
		System.out.println("ALICE: Initialization ...");
		KeyAgreement aliceKeyAgree = KeyAgreement.getInstance("DH");
		aliceKeyAgree.init(aliceKpair.getPrivate());

		// Alice encodes her public key, and sends it over to Bob.
		byte[] alicePubKeyEnc = aliceKpair.getPublic().getEncoded();

		/*
		 * Let's turn over to Bob. Bob has received Alice's public key in
		 * encoded format. He instantiates a DH public key from the encoded key
		 * material.
		 */
		KeyFactory bobKeyFac = KeyFactory.getInstance("DH");
		X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(alicePubKeyEnc);
		PublicKey alicePubKey = bobKeyFac.generatePublic(x509KeySpec);

		/*
		 * Bob gets the DH parameters associated with Alice's public key. He
		 * must use the same parameters when he generates his own key pair.
		 */
		DHParameterSpec dhParamSpec = ((DHPublicKey) alicePubKey).getParams();

		// Bob creates his own DH key pair
		System.out.println("BOB: Generate DH keypair ...");
		KeyPairGenerator bobKpairGen = KeyPairGenerator.getInstance("DH");
		bobKpairGen.initialize(dhParamSpec);
		KeyPair bobKpair = bobKpairGen.generateKeyPair();

		// Bob creates and initializes his DH KeyAgreement object
		System.out.println("BOB: Initialization ...");
		KeyAgreement bobKeyAgree = KeyAgreement.getInstance("DH");
		bobKeyAgree.init(bobKpair.getPrivate());

		// Bob encodes his public key, and sends it over to Alice.
		byte[] bobPubKeyEnc = bobKpair.getPublic().getEncoded();

		/*
		 * Alice uses Bob's public key for the first (and only) phase of her
		 * version of the DH protocol. Before she can do so, she has to
		 * instantiate a DH public key from Bob's encoded key material.
		 */
		KeyFactory aliceKeyFac = KeyFactory.getInstance("DH");
		x509KeySpec = new X509EncodedKeySpec(bobPubKeyEnc);
		PublicKey bobPubKey = aliceKeyFac.generatePublic(x509KeySpec);
		System.out.println("ALICE: Execute PHASE1 ...");
		aliceKeyAgree.doPhase(bobPubKey, true);

		/*
		 * Bob uses Alice's public key for the first (and only) phase of his
		 * version of the DH protocol.
		 */
		System.out.println("BOB: Execute PHASE1 ...");
		bobKeyAgree.doPhase(alicePubKey, true);

		/*
		 * At this stage, both Alice and Bob have completed the DH key agreement
		 * protocol. Both generate the (same) shared secret.
		 */
		byte[] aliceSharedSecret = aliceKeyAgree.generateSecret();
		int aliceLen = aliceSharedSecret.length;

		byte[] bobSharedSecret = new byte[aliceLen];
		int bobLen;
		try {
			// show example of what happens if you
			// provide an output buffer that is too short
			bobLen = bobKeyAgree.generateSecret(bobSharedSecret, 1);
		} catch (ShortBufferException e) {
			System.out.println(e.getMessage());
		}
		// provide output buffer of required size
		bobLen = bobKeyAgree.generateSecret(bobSharedSecret, 0);

		System.out.println("Alice secret: " + toHexString(aliceSharedSecret));
		System.out.println("Bob secret: " + toHexString(bobSharedSecret));

		if (!java.util.Arrays.equals(aliceSharedSecret, bobSharedSecret))
			throw new Exception("Shared secrets differ");
		System.out.println("Shared secrets are the same");

		/*
		 * Now let's return the shared secret as a SecretKey object and use it
		 * for encryption. First, we generate SecretKeys for the "DES" algorithm
		 * (based on the raw shared secret data) and then we use DES in ECB mode
		 * as the encryption algorithm. DES in ECB mode does not require any
		 * parameters.
		 * 
		 * Then we use DES in CBC mode, which requires an initialization vector
		 * (IV) parameter. In CBC mode, you need to initialize the Cipher object
		 * with an IV, which can be supplied using the
		 * javax.crypto.spec.IvParameterSpec class. Note that you have to use
		 * the same IV for encryption and decryption: If you use a different IV
		 * for decryption than you used for encryption, decryption will fail.
		 * 
		 * NOTE: If you do not specify an IV when you initialize the Cipher
		 * object for encryption, the underlying implementation will generate a
		 * random one, which you have to retrieve using the
		 * javax.crypto.Cipher.getParameters() method, which returns an instance
		 * of java.security.AlgorithmParameters. You need to transfer the
		 * contents of that object (e.g., in encoded format, obtained via the
		 * AlgorithmParameters.getEncoded() method) to the party who will do the
		 * decryption. When initializing the Cipher for decryption, the
		 * (reinstantiated) AlgorithmParameters object must be passed to the
		 * Cipher.init() method.
		 */
		System.out.println("Return shared secret as SecretKey object ...");
		// Bob
		// NOTE: The call to bobKeyAgree.generateSecret above reset the key
		// agreement object, so we call doPhase again prior to another
		// generateSecret call
		bobKeyAgree.doPhase(alicePubKey, true);
		SecretKey bobDesKey = bobKeyAgree.generateSecret("DES");

		// Alice
		// NOTE: The call to aliceKeyAgree.generateSecret above reset the key
		// agreement object, so we call doPhase again prior to another
		// generateSecret call
		aliceKeyAgree.doPhase(bobPubKey, true);
		SecretKey aliceDesKey = aliceKeyAgree.generateSecret("DES");

		/*
		 * Bob encrypts, using DES in ECB mode
		 */
		Cipher bobCipher = Cipher.getInstance("DES/ECB/PKCS5Padding");
		bobCipher.init(Cipher.ENCRYPT_MODE, bobDesKey);

		byte[] cleartext = "This is just an example".getBytes();
		byte[] ciphertext = bobCipher.doFinal(cleartext);

		/*
		 * Alice decrypts, using DES in ECB mode
		 */
		Cipher aliceCipher = Cipher.getInstance("DES/ECB/PKCS5Padding");
		aliceCipher.init(Cipher.DECRYPT_MODE, aliceDesKey);
		byte[] recovered = aliceCipher.doFinal(ciphertext);

		if (!java.util.Arrays.equals(cleartext, recovered))
			throw new Exception("DES in CBC mode recovered text is "
					+ "different from cleartext");
		System.out.println("DES in ECB mode recovered text is "
				+ "same as cleartext");

		/*
		 * Bob encrypts, using DES in CBC mode
		 */
		bobCipher = Cipher.getInstance("DES/CBC/PKCS5Padding");
		bobCipher.init(Cipher.ENCRYPT_MODE, bobDesKey);

		cleartext = "This is just an example".getBytes();
		ciphertext = bobCipher.doFinal(cleartext);
		// Retrieve the parameter that was used, and transfer it to Alice in
		// encoded format
		byte[] encodedParams = bobCipher.getParameters().getEncoded();

		/*
		 * Alice decrypts, using DES in CBC mode
		 */
		// Instantiate AlgorithmParameters object from parameter encoding
		// obtained from Bob
		AlgorithmParameters params = AlgorithmParameters.getInstance("DES");
		params.init(encodedParams);
		aliceCipher = Cipher.getInstance("DES/CBC/PKCS5Padding");
		aliceCipher.init(Cipher.DECRYPT_MODE, aliceDesKey, params);
		recovered = aliceCipher.doFinal(ciphertext);

		if (!java.util.Arrays.equals(cleartext, recovered))
			throw new Exception("DES in CBC mode recovered text is "
					+ "different from cleartext");
		System.out.println("DES in CBC mode recovered text is "
				+ "same as cleartext");
	}

	/*
	 * Converts a byte to hex digit and writes to the supplied buffer
	 */
	private void byte2hex(byte b, StringBuffer buf) {
		char[] hexChars = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
				'A', 'B', 'C', 'D', 'E', 'F' };
		int high = ((b & 0xf0) >> 4);
		int low = (b & 0x0f);
		buf.append(hexChars[high]);
		buf.append(hexChars[low]);
	}

	/*
	 * Converts a byte array to hex string
	 */
	private String toHexString(byte[] block) {
		StringBuffer buf = new StringBuffer();

		int len = block.length;

		for (int i = 0; i < len; i++) {
			byte2hex(block[i], buf);
			if (i < len - 1) {
				buf.append(":");
			}
		}
		return buf.toString();
	}

	/*
	 * Prints the usage of this test.
	 */
	private void usage() {
		System.err.print("DHKeyAgreement usage: ");
		System.err.println("[-gen]");
	}

	// The 1024 bit Diffie-Hellman modulus values used by SKIP
	private static final byte skip1024ModulusBytes[] = { (byte) 0xF4,
			(byte) 0x88, (byte) 0xFD, (byte) 0x58, (byte) 0x4E, (byte) 0x49,
			(byte) 0xDB, (byte) 0xCD, (byte) 0x20, (byte) 0xB4, (byte) 0x9D,
			(byte) 0xE4, (byte) 0x91, (byte) 0x07, (byte) 0x36, (byte) 0x6B,
			(byte) 0x33, (byte) 0x6C, (byte) 0x38, (byte) 0x0D, (byte) 0x45,
			(byte) 0x1D, (byte) 0x0F, (byte) 0x7C, (byte) 0x88, (byte) 0xB3,
			(byte) 0x1C, (byte) 0x7C, (byte) 0x5B, (byte) 0x2D, (byte) 0x8E,
			(byte) 0xF6, (byte) 0xF3, (byte) 0xC9, (byte) 0x23, (byte) 0xC0,
			(byte) 0x43, (byte) 0xF0, (byte) 0xA5, (byte) 0x5B, (byte) 0x18,
			(byte) 0x8D, (byte) 0x8E, (byte) 0xBB, (byte) 0x55, (byte) 0x8C,
			(byte) 0xB8, (byte) 0x5D, (byte) 0x38, (byte) 0xD3, (byte) 0x34,
			(byte) 0xFD, (byte) 0x7C, (byte) 0x17, (byte) 0x57, (byte) 0x43,
			(byte) 0xA3, (byte) 0x1D, (byte) 0x18, (byte) 0x6C, (byte) 0xDE,
			(byte) 0x33, (byte) 0x21, (byte) 0x2C, (byte) 0xB5, (byte) 0x2A,
			(byte) 0xFF, (byte) 0x3C, (byte) 0xE1, (byte) 0xB1, (byte) 0x29,
			(byte) 0x40, (byte) 0x18, (byte) 0x11, (byte) 0x8D, (byte) 0x7C,
			(byte) 0x84, (byte) 0xA7, (byte) 0x0A, (byte) 0x72, (byte) 0xD6,
			(byte) 0x86, (byte) 0xC4, (byte) 0x03, (byte) 0x19, (byte) 0xC8,
			(byte) 0x07, (byte) 0x29, (byte) 0x7A, (byte) 0xCA, (byte) 0x95,
			(byte) 0x0C, (byte) 0xD9, (byte) 0x96, (byte) 0x9F, (byte) 0xAB,
			(byte) 0xD0, (byte) 0x0A, (byte) 0x50, (byte) 0x9B, (byte) 0x02,
			(byte) 0x46, (byte) 0xD3, (byte) 0x08, (byte) 0x3D, (byte) 0x66,
			(byte) 0xA4, (byte) 0x5D, (byte) 0x41, (byte) 0x9F, (byte) 0x9C,
			(byte) 0x7C, (byte) 0xBD, (byte) 0x89, (byte) 0x4B, (byte) 0x22,
			(byte) 0x19, (byte) 0x26, (byte) 0xBA, (byte) 0xAB, (byte) 0xA2,
			(byte) 0x5E, (byte) 0xC3, (byte) 0x55, (byte) 0xE9, (byte) 0x2F,
			(byte) 0x78, (byte) 0xC7 };

	// The SKIP 1024 bit modulus
	private static final BigInteger skip1024Modulus = new BigInteger(1,
			skip1024ModulusBytes);

	// The base used with the SKIP 1024 bit modulus
	private static final BigInteger skip1024Base = BigInteger.valueOf(2);
}